package tree

// https://leetcode-cn.com/problems/implement-trie-prefix-tree/
// https://blog.csdn.net/weixin_39778570/article/details/81990417
// 前缀树, 解法: 字典树
// 用途：引擎搜索
// 每一层保留 26 个英文字母对应的所有下标，可能会导致空间利用比较高
type TrieNode struct {
	isWord    bool
	charSlice []*TrieNode
}

func NewTrieNode() *TrieNode {
	currentTrieNode := new(TrieNode)
	currentTrieNode.charSlice = make([]*TrieNode, 26)
	return currentTrieNode
}

func (node *TrieNode) HasNextChar(char int32) bool {
	return nil != node.charSlice[char-'a']
}

func (node *TrieNode) GetNextChar(char int32) *TrieNode {
	return node.charSlice[char-'a']
}

func (node *TrieNode) NewNextChar(char int32) *TrieNode {
	node.charSlice[char-'a'] = NewTrieNode()
	return node.charSlice[char-'a']
}

func (node *TrieNode) SetIsWord() {
	node.isWord = true
}

type Trie struct {
	node *TrieNode
}

/** Initialize your data structure here. */
func Constructor() Trie {
	currentTrie := Trie{}
	currentTrie.node = NewTrieNode()
	return currentTrie
}

/** Inserts a word into the trie. */
func (this *Trie) Insert(word string) {
	currentNode := this.node
	for _, currentChar := range word {
		if currentNode.HasNextChar(currentChar) {
			currentNode = currentNode.GetNextChar(currentChar)
		} else {
			currentNode = currentNode.NewNextChar(currentChar)
		}
	}
	// 记得最后一个位置要标记成是单词
	currentNode.SetIsWord()
}

/** Returns if the word is in the trie. */
// 查找时间复杂度: O(len(word))
func (this *Trie) Search(word string) bool {
	currentNode := this.node
	for _, currentChar := range word {
		if nil != currentNode.charSlice[currentChar-'a'] {
			currentNode = currentNode.charSlice[currentChar-'a']
		} else {
			return false
		}
	}
	if currentNode.isWord {
		return true
	} else {
		return false
	}
}

/** Returns if there is any word in the trie that starts with the given prefix. */
func (this *Trie) StartsWith(prefix string) bool {
	currentNode := this.node
	for _, currentChar := range prefix {
		if nil != currentNode.charSlice[currentChar-'a'] {
			currentNode = currentNode.charSlice[currentChar-'a']
		} else {
			return false
		}
	}
	return true
}

/**
 * Your Trie object will be instantiated and called as such:
 * obj := Constructor();
 * obj.Insert(word);
 * param_2 := obj.Search(word);
 * param_3 := obj.StartsWith(prefix);
 */
